import sqlite3
import pandas as pd
import numpy as np
from collections import Counter
import random
from datetime import datetime
import matplotlib.pyplot as plt

class DLTPredictor:
    def __init__(self):
        print("[DEBUG] 正在连接数据库 lottery.db...")
        try:
            self.conn = sqlite3.connect('lottery.db')
            print("[DEBUG] 数据库连接成功")
        except Exception as e:
            print(f"[ERROR] 数据库连接失败: {e}")
            raise
            
        self.recent_draws = 100  # 分析最近100期数据
        self.front_balls_range = range(1, 36)  # 前区范围1-35
        self.back_balls_range = range(1, 13)  # 后区范围1-12
        print(f"[DEBUG] 初始化完成，分析最近 {self.recent_draws} 期数据")
        
    def get_historical_data(self):
        """获取历史开奖数据"""
        print(f"[DEBUG] 正在查询最近 {self.recent_draws} 期开奖数据...")
        query = """
            SELECT draw_num, draw_date, 
                   red_ball_1, red_ball_2, red_ball_3, red_ball_4, red_ball_5, 
                   blue_ball_1, blue_ball_2
            FROM dlt_draws
            ORDER BY draw_num DESC
            LIMIT ?
        """
        
        try:
            df = pd.read_sql_query(query, self.conn, params=(self.recent_draws,))
            print(f"[DEBUG] 成功获取 {len(df)} 条开奖记录")
            if len(df) == 0:
                print("[WARNING] 查询结果为空，请检查数据库是否有数据")
            return df
        except Exception as e:
            print(f"[ERROR] 数据查询失败: {e}")
            raise
    
    def analyze_frequency(self, df):
        """分析号码出现频率"""
        # 前区频率分析
        front_balls = []
        for i in range(1, 6):
            front_balls.extend(df[f'red_ball_{i}'].tolist())
        
        front_freq = Counter(front_balls)
        for num in self.front_balls_range:
            if num not in front_freq:
                front_freq[num] = 0
        
        # 后区频率分析
        back_freq = Counter()
        for i in range(1, 3):
            back_freq.update(df[f'blue_ball_{i}'].tolist())
        
        for num in self.back_balls_range:
            if num not in back_freq:
                back_freq[num] = 0
        
        return front_freq, back_freq
    
    def analyze_hot_cold(self, front_freq, back_freq):
        """分析冷热号"""
        # 按频率排序
        sorted_front = sorted(front_freq.items(), key=lambda x: x[1], reverse=True)
        sorted_back = sorted(back_freq.items(), key=lambda x: x[1], reverse=True)
        
        # 热号(出现频率前1/3)
        hot_front = [num for num, _ in sorted_front[:12]]
        hot_back = [num for num, _ in sorted_back[:4]]
        
        # 温号(出现频率中间1/3)
        warm_front = [num for num, _ in sorted_front[12:24]]
        warm_back = [num for num, _ in sorted_back[4:8]]
        
        # 冷号(出现频率后1/3)
        cold_front = [num for num, _ in sorted_front[24:]]
        cold_back = [num for num, _ in sorted_back[8:]]
        
        return {
            'hot_front': hot_front,
            'warm_front': warm_front,
            'cold_front': cold_front,
            'hot_back': hot_back,
            'warm_back': warm_back,
            'cold_back': cold_back
        }
    
    def generate_predictions(self, count=10):
        """生成预测号码"""
        print("[DEBUG] 开始生成预测号码...")
        print(f"[DEBUG] 将生成 {count} 组预测号码")
        
        df = self.get_historical_data()
        print("[DEBUG] 正在分析号码频率...")
        front_freq, back_freq = self.analyze_frequency(df)
        
        print("[DEBUG] 正在分析冷热号...")
        hot_cold_analysis = self.analyze_hot_cold(front_freq, back_freq)
        
        # 生成预测号码
        predictions = []
        for _ in range(count):
            # 前区号码选择
            front_candidates = set()
            
            # 加入一些热号
            front_candidates.update(random.sample(hot_cold_analysis['hot_front'], k=min(2, len(hot_cold_analysis['hot_front']))))
            
            # 加入一些温号
            front_candidates.update(random.sample(hot_cold_analysis['warm_front'], k=min(2, len(hot_cold_analysis['warm_front']))))
            
            # 如果不足5个，从剩余号码中随机选择
            remaining_count = 5 - len(front_candidates)
            if remaining_count > 0:
                remaining_numbers = list(set(self.front_balls_range) - front_candidates)
                front_candidates.update(random.sample(remaining_numbers, k=remaining_count))
            
            # 确保只有5个前区号码
            if len(front_candidates) > 5:
                front_candidates = set(random.sample(list(front_candidates), k=5))
            
            # 后区号码选择
            back_candidates = []
            
            # 有70%的概率选择热门后区号码
            if random.random() < 0.7:
                back_candidates.extend(hot_cold_analysis['hot_back'])
            else:
                back_candidates.extend(hot_cold_analysis['warm_back'])
            
            # 如果没有候选后区号码，从所有后区号码中随机选择
            if not back_candidates:
                back_candidates = list(self.back_balls_range)
            
            back_balls = random.sample(back_candidates, k=2)
            
            # 添加到预测结果
            predictions.append((sorted(list(front_candidates)), sorted(back_balls)))
        
        print("\n=== 推荐号码组合 ===\n")
        for i, (front_balls, back_balls) in enumerate(predictions, 1):
            print(f"第{i}组: {', '.join(str(num) for num in front_balls)} + {', '.join(str(num) for num in back_balls)}")
        
        return predictions

if __name__ == "__main__":
    predictor = DLTPredictor()
    predictor.generate_predictions(5)